#define BLOCK_SIZE 128

#define G 6.67384e-11f
#define EPS 1e-6f // minimum distance squared

#include "octree.h"

int offsets[] = {0, 1, 9, 73, 585, 4681, 37449, 299593, 2396745, 19173961};
 __device__ __constant__ int d_offsets[] = {0, 1, 9, 73, 585, 4681, 37449, 299593, 2396745, 19173961};

int getDepth(int x)
{
	int depth = 0;

	while (x >= offsets[depth])
		++depth;

	return depth - 1;
}

void printBinary(int x)
{
	for (int i = 31; i >= 0; --i)
	{
		std::cout << ((x >> i) & 1);
	}
}

__global__ void scanMin(float4* data, float4* totals, int n)
{
	//extern __shared__ int temp[];
	__shared__ float4 temp[BLOCK_SIZE];

	int block_offset = blockIdx.x*blockDim.x;
	int t_id = threadIdx.x;
	int x = block_offset + t_id;

	// if n is not a power of 2, pad local memory with 0
	if (x < n)
		temp[t_id] = data[x];
	else
		temp[t_id] = make_float4(1.0f/0.0f);

	int offset = 1;

	for (int d = blockDim.x >> 1; d > 0; d >>= 1)
	{
		__syncthreads();

		if (t_id < d)
		{
			int a = offset*(2*t_id + 1) - 1;
			int b = a + offset;

			//temp[b] += temp[a];
			temp[b].x = temp[b].x < temp[a].x ? temp[b].x : temp[a].x;
			temp[b].y = temp[b].y < temp[a].y ? temp[b].y : temp[a].y;
			temp[b].z = temp[b].z < temp[a].z ? temp[b].z : temp[a].z;
		}

		offset *= 2;
	}

	if (t_id == 0)
		totals[blockIdx.x] = temp[blockDim.x - 1];

	__syncthreads();

	if (x < n)
		data[x] = temp[t_id];
}

__global__ void scanMax(float4* data, float4* totals, int n)
{
	//extern __shared__ int temp[];
	__shared__ float4 temp[BLOCK_SIZE];

	int block_offset = blockIdx.x*blockDim.x;
	int t_id = threadIdx.x;
	int x = block_offset + t_id;

	// if n is not a power of 2, pad local memory with 0
	if (x < n)
		temp[t_id] = data[x];
	else
		temp[t_id] = make_float4(-1.0f/0.0f);

	int offset = 1;

	for (int d = blockDim.x >> 1; d > 0; d >>= 1)
	{
		__syncthreads();

		if (t_id < d)
		{
			int a = offset*(2*t_id + 1) - 1;
			int b = a + offset;

			//temp[b] += temp[a];
			temp[b].x = temp[b].x > temp[a].x ? temp[b].x : temp[a].x;
			temp[b].y = temp[b].y > temp[a].y ? temp[b].y : temp[a].y;
			temp[b].z = temp[b].z > temp[a].z ? temp[b].z : temp[a].z;
		}

		offset *= 2;
	}

	if (t_id == 0)
		totals[blockIdx.x] = temp[blockDim.x - 1];

	__syncthreads();

	if (x < n)
		data[x] = temp[t_id];
}

float4 getMin(float4* d_data, int n)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((n-1)/dimb.x + 1, 1);

	float4* d_mins;
	cudaMalloc(&d_mins, dimg.x*sizeof(float4));

	scanMin<<<dimg, dimb>>>(d_data, d_mins, n);

	float4 min;

	if (dimg.x > 1)
		min = getMin(d_mins, dimg.x);
	else
		cudaMemcpy(&min, d_mins, sizeof(float4), cudaMemcpyDeviceToHost);

	cudaFree(d_mins);

	return min;
}

float4 getMax(float4* d_data, int n)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((n-1)/dimb.x + 1, 1);

	float4* d_maxs;
	cudaMalloc(&d_maxs, dimg.x*sizeof(float4));

	scanMax<<<dimg, dimb>>>(d_data, d_maxs, n);

	float4 max;

	if (dimg.x > 1)
		max = getMax(d_maxs, dimg.x);
	else
		cudaMemcpy(&max, d_maxs, sizeof(float4), cudaMemcpyDeviceToHost);

	cudaFree(d_maxs);

	return max;
}

void getBB(float4* d_pos, int n, float3& min, float3& max)
{
	float4* temp;
	cudaMalloc(&temp, n*sizeof(float4));

	cudaMemcpy(temp, d_pos, n*sizeof(float4), cudaMemcpyDeviceToDevice);
	min = make_float3(getMin(temp, n));

	cudaMemcpy(temp, d_pos, n*sizeof(float4), cudaMemcpyDeviceToDevice);
	max = make_float3(getMax(temp, n));

	cudaFree(temp);
}

__host__ __device__ int spread3(int x)
{
	x = (x | (x << 16)) & 0x030000FF;
	x = (x | (x <<  8)) & 0x0300F00F;
	x = (x | (x <<  4)) & 0x030C30C3;
	x = (x | (x <<  2)) & 0x09249249;

	return x;
}

__host__ __device__ int3 getCoords(float3 pos, float3 bb_min, float bb_len, int divs)
{
	float3 f_coords = (pos - bb_min)/(make_float3(bb_len/divs));
	return clamp(make_int3(f_coords.x, f_coords.y, f_coords.z), 0, divs - 1);
}

__host__ __device__ int getMortonCode(float3 pos, float3 bb_min, float bb_len, int divs)
{
	int3 coords = getCoords(pos, bb_min, bb_len, divs);
	return (spread3(coords.x) << 2) | (spread3(coords.y) << 1) | spread3(coords.z);
}

__global__ void setIndices(record_type* records, float4* pos, int n,
	float3 bb_min, float bb_len, int divs)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
	{
		records[x].first = getMortonCode(make_float3(pos[x]), bb_min, bb_len, divs);
		records[x].second = x;
	}
}

__global__ void initNodeLevel(Node<float4>* nodes, int n, int depth)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
	{
		int curr_node = d_offsets[depth] + x;
		int parent_node = d_offsets[depth-1] + (x >> 3);

		// w component stores the half length of current node
		float4 parent_centre = nodes[parent_node].centre;
		float len = parent_centre.w;
		float half_len = len/2.0f;
		float4 min_centre = parent_centre - make_float4(half_len);
		float4 offset = make_float4(len*make_float3((x >> 2) & 1, (x >> 1) & 1, (x & 1)));
		
		nodes[curr_node].centre = min_centre + offset;
		nodes[curr_node].children = 8*curr_node + 1;
		nodes[curr_node].num_children = 0;
		nodes[curr_node].item = make_float4(0.0f);
	}
}

__global__ void setLeaves(Node<float4>* nodes, record_type* records, float4* pos, int n, int depth)
{
	__shared__ record_type temp[BLOCK_SIZE + 1];

	int x = blockIdx.x*blockDim.x + threadIdx.x;
	int t_id = threadIdx.x;
	int offset = d_offsets[depth];

	if (x < n)
	{
		temp[t_id + 1] = records[x];

		if (t_id == 0)
		{
			if (x != 0)
				temp[0] = records[x - 1];
		}
	}

	__syncthreads();

	if (x < n)
	{
		if (temp[t_id + 1].first != temp[t_id].first || x == 0)
		{
			// this is a beginning of new node
			nodes[offset + temp[t_id + 1].first].children = x;

			// find average of all bodies in this node
			float4 pos_avg = make_float4(0.0f);
			int curr_body = x;
			int bin = records[curr_body].first;
			while (curr_body < n && records[curr_body].first == bin)
			{
				pos_avg += pos[records[curr_body].second];
				++curr_body;
			}

			// set leaf child pointer, number of children and average of bodies
			int num_children = curr_body - x;
			pos_avg.x /= num_children;
			pos_avg.y /= num_children;
			pos_avg.z /= num_children;
			nodes[offset + temp[t_id + 1].first].num_children = num_children;
			nodes[offset + temp[t_id + 1].first].item = pos_avg;
		}
	}
}

__global__ void updateOctreeLevel(Node<float4>* nodes, int n, int depth)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;
	int offset = d_offsets[depth];

	if (x < n)
	{
		int curr_children = nodes[offset + x].children;

		int total_children = 0;
		float4 avg_pos = make_float4(0.0f);

		for (int i = 0; i != 8; ++i)
		{
			int num_children = nodes[curr_children + i].num_children;
			total_children += num_children;

			float4 curr_pos = nodes[curr_children + i].item;
			curr_pos.x *= num_children;
			curr_pos.y *= num_children;
			curr_pos.z *= num_children;
			avg_pos += curr_pos;
		}

		if (total_children != 0)
		{
			avg_pos.x /= total_children;
			avg_pos.y /= total_children;
			avg_pos.z /= total_children;
		}

		nodes[offset + x].num_children = total_children;
		nodes[offset + x].item = avg_pos;
	}
}

//calclulate the acceleration of a to b
inline __host__ __device__ float3 calcGravityAccel(float4* pos_a, float4* pos_b)
{
    float3 r = make_float3(*pos_a) - make_float3(*pos_b);   // 3 flops
    float inv_dis = rsqrtf(dot(r, r) + EPS);
    float inv_dis_3 = inv_dis*inv_dis*inv_dis;

    return (-G*(pos_b->w)*inv_dis_3)*r;                 // 5 flops
}

__global__ void updateVel(Node<float4>* nodes, record_type* records, float4* pos, float4* vel, int n, int depth, float dt)
{
	int x = blockIdx.x*blockDim.x + threadIdx.x;

	if (x < n)
	{
		float4 pos_a = pos[x];
		float3 acc = make_float3(0.0f);

		int curr_depth = 0;
		int curr_node = 0;
		do
		{
			//vel[x].z += 1;
			bool expanded = false;

			Node<float4> node = nodes[d_offsets[curr_depth] + curr_node];

			if (node.num_children != 0)
			{
				float3 diff = make_float3(pos_a) - make_float3(node.centre);
				float d = sqrtf(dot(diff, diff));

				if (2.0f*node.centre.w/d < 1.0f)
				{
					// only add node average
					acc += calcGravityAccel(&pos_a, &node.item);
				}
				else if (curr_depth == depth)
				{
					// average all in leaf node
					int curr_body = node.children;
					int bin = records[curr_body].first;

					while (curr_body < n && records[curr_body].first == bin)
					{
						float4 pos_b = pos[records[curr_body].second];
						acc += calcGravityAccel(&pos_a, &pos_b);
						++curr_body;
					}
				}
				else
				{
					// expand node
					++curr_depth;
					curr_node <<= 3;
					expanded = true;
				}
			}

			if (!expanded)
			{
				while ((curr_node & 7) == 7)
				{
					--curr_depth;
					curr_node >>= 3;
				}

				++curr_node;
			}
		} while (curr_depth != 0);

		vel[x] += make_float4(dt*acc);
	}
}

Octree::Octree(int _n, int _depth) :
	n(_n),
	depth(_depth)
{
	cudaMalloc(&d_nodes, offsets[depth+1]*sizeof(Node<float4>));
	cudaMalloc(&d_records, n*sizeof(record_type));
}

Octree::~Octree()
{
	cudaFree(d_nodes);
	cudaFree(d_records);
}

void Octree::update(float4* d_pos)
{
	float3 min, max;
	getBB(d_pos, n, min, max);

	float3 centre = (min + max)/2.0f;
	float3 diff = max - min;
	float len = fmaxf(diff.x, fmaxf(diff.y, diff.z));
	float half_len = len/2.0f;
	float4 bb = make_float4(centre, half_len);

	// initialize root node
	Node<float4> root_node;
	root_node.centre = make_float4(centre, half_len);
	root_node.children = 1;
	root_node.num_children = 0;
	root_node.item = make_float4(0.0f);

	cudaMemcpy(d_nodes, &root_node, sizeof(Node<float4>), cudaMemcpyHostToDevice);

	// initialize octree nodes starting from root
	for (int i = 1; i != depth + 1; ++i)
	{
		dim3 dimb(BLOCK_SIZE, 1);
		dim3 dimg((pow(8, i)-1)/dimb.x + 1, 1);
		initNodeLevel<<<dimg, dimb>>>(d_nodes, pow(8, i), i);
	}

	// get morton number indices, sort and link bins to octree leaves
	dim3 dimb_n(BLOCK_SIZE, 1);
	dim3 dimg_n((n-1)/dimb_n.x + 1, 1);
	setIndices<<<dimg_n, dimb_n>>>(d_records, d_pos, n,
		centre - make_float3(half_len), len, pow(2, depth));

	radixSort(d_records, n);

	setLeaves<<<dimg_n, dimb_n>>>(d_nodes, d_records, d_pos, n, depth);

	// update octree averages
	for (int i = depth - 1; i >= 0; --i)
	{
		dim3 dimb(BLOCK_SIZE, 1);
		dim3 dimg((pow(8, i)-1)/dimb.x + 1, 1);
		updateOctreeLevel<<<dimg, dimb>>>(d_nodes, pow(8, i), i);
	}
}

void Octree::integrate(float4* d_pos, float4* d_vel, float dt)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((n-1)/dimb.x + 1, 1);
	updateVel<<<dimg, dimb>>>(d_nodes, d_records, d_pos, d_vel, n, depth, dt);
}

void Octree::print(std::vector<float4>& pos)
{
	std::vector<record_type> records(n);

	cudaMemcpy(&records[0], d_records, n*sizeof(record_type), cudaMemcpyDeviceToHost);

	for (int i = 0; i != records.size(); ++i)
	{
		int curr_body = records[i].second;
		std::cout << records[i].first << " " << records[i].second << " "
			<< "(" << pos[curr_body].x << " " << pos[curr_body].y << " "
			<< pos[curr_body].z << " " << pos[curr_body].w << ")" << " "
			<< std::endl;
	}

	std::cout << std::endl;

	std::vector<Node<float4> > nodes(offsets[depth+1]);

	cudaMemcpy(&nodes[0], d_nodes, nodes.size()*sizeof(Node<float4>), cudaMemcpyDeviceToHost);

	//for (int i = 0; i != nodes.size(); ++i)
	for (int i = 0; i != offsets[3]; ++i)
	{
		int depth_i = i - offsets[getDepth(i)];
		int mort = getMortonCode(make_float3(nodes[i].centre), make_float3(-1, -1, -1), 2, 16);
		printBinary(depth_i);
		std::cout << " " << getDepth(i) << " " << i << " " << depth_i << " " << (depth_i >> 3) + offsets[getDepth(i) - 1] << " "
			<< nodes[i].children << " " << nodes[i].num_children << " "
			<< "(" << nodes[i].centre.x << " " << nodes[i].centre.y << " "
			<< nodes[i].centre.z << " " << nodes[i].centre.w << ")" << " "
			<< "(" << nodes[i].item.x << " " << nodes[i].item.y << " "
			<< nodes[i].item.z << " " << nodes[i].item.w << ")"
			<< std::endl;
	}

	std::cout << std::endl;
}
